import numpy as np
import torch
import torch.nn as nn


def event2img_cover(event,
                    channel = 3, # 不要改动
                    num_frame=1,
                    needed= -1, 
                    
                    w=1280, 
                    h=720, 
                    device='cuda' if torch.cuda.is_available() else 'cpu'):
    """
    Args:
        event: 事件数据，包含'x','y','p','t'字段
        num_frame: 帧数
        w: 图像宽度
        h: 图像高度
        device: 计算设备 ('cuda' 或 'cpu')
    
    Returns:
        单帧或多帧图像
    """
    
    
    if num_frame == 1:
        # 创建图像张量 [H, W, 3]
        image = torch.zeros((h, w, channel), dtype=torch.uint8, device=device)
        
        # 提取事件数据并转移到GPU
        x_coords = torch.tensor([e['x'] for e in event], dtype=torch.long, device=device)
        y_coords = torch.tensor([e['y'] for e in event], dtype=torch.long, device=device)
        polarities = torch.tensor([e['p'] for e in event], dtype=torch.int32, device=device)
        
        # 按时间顺序处理事件（新事件覆盖旧事件）
        # 由于事件数据通常已经按时间排序，我们按顺序处理即可
        for i in range(len(event)):
            if polarities[i] == 1:
                image[y_coords[i], x_coords[i]] = torch.tensor([200, 200, 200], dtype=torch.uint8, device=device)
            else:
                image[y_coords[i], x_coords[i]] = torch.tensor([100, 100, 100], dtype=torch.uint8, device=device)
        
        # 将结果传回CPU并转换为numpy
        return image.cpu().numpy()
    
    else:

        # 多帧情况
        t_min, t_max = event[0]['t'], event[-1]['t']
        t_bins = torch.linspace(t_min, t_max, num_frame + 1, device=device)
        
        
        
        # 提取所有事件数据
        x_coords = torch.tensor([e['x'] for e in event], dtype=torch.long, device=device)
        y_coords = torch.tensor([e['y'] for e in event], dtype=torch.long, device=device)
        polarities = torch.tensor([e['p'] for e in event], dtype=torch.int32, device=device)
        timestamps = torch.tensor([e['t'] for e in event], dtype=torch.float32, device=device)


        
        # 为每个事件计算帧索引
        frame_indices = torch.bucketize(timestamps, t_bins, right=True) - 1
        frame_indices = torch.clamp(frame_indices, 0, num_frame - 1)

        if needed == -1:
            image = torch.zeros(( h, w, channel), dtype=torch.uint8, device=device)

            frame_mask = frame_indices == num_frame-1
            frame_events_count = frame_mask.sum().item()
            
            if frame_events_count == 0:
                return image.cpu().numpy()
                
            # 提取该帧的事件数据
            frame_x = x_coords[frame_mask]
            frame_y = y_coords[frame_mask]
            frame_p = polarities[frame_mask]
            
            # 按时间顺序处理该帧内的事件（新事件覆盖旧事件）
            for i in range(frame_events_count):
                if polarities[i] == 1:
                    image[y_coords[i], x_coords[i]] = torch.tensor([200, 200, 200], dtype=torch.uint8, device=device)
                else:
                    image[y_coords[i], x_coords[i]] = torch.tensor([100, 100, 100], dtype=torch.uint8, device=device)

            # 将结果传回CPU并转换为numpy
            return image.cpu().numpy()
        
        else:    

            # 创建多帧图像张量 [num_frame, H, W, 3]
            images = torch.zeros((num_frame, h, w, channel), dtype=torch.uint8, device=device)
            
            # 按时间顺序处理每个帧内的事件
            for frame_idx in range(num_frame):
                # 获取该帧的所有事件
                frame_mask = frame_indices == frame_idx
                frame_events_count = frame_mask.sum().item()
                
                if frame_events_count == 0:
                    continue
                    
                # 提取该帧的事件数据
                frame_x = x_coords[frame_mask]
                frame_y = y_coords[frame_mask]
                frame_p = polarities[frame_mask]
                
                # 按时间顺序处理该帧内的事件（新事件覆盖旧事件）
                for i in range(frame_events_count):
                    if frame_p[i] == 1:
                        images[frame_idx, frame_y[i], frame_x[i]] = torch.tensor([200, 200, 200], dtype=torch.uint8, device=device)
                    else:
                        images[frame_idx, frame_y[i], frame_x[i]] = torch.tensor([100, 100, 100], dtype=torch.uint8, device=device)
            
            # 将结果传回CPU
            images_np = images.cpu().numpy()
            return [images_np[i] for i in range(num_frame)]



def event2img_count(event, 
                    channel = 3,
                    num_frame=1,
                    needed= -1, 
                    threshold = 1, 
                    w=1280, 
                    h=720, 
                    device='cuda' if torch.cuda.is_available() else 'cpu'):
    """
    将事件数据转换为计数图像
    - 初始值为128
    - 正事件: +1
    - 负事件: -1
    - 限制在[0, 255]范围内
    - 根据阈值输出图像,默认为1
    """
    
    # 将事件数据转移到GPU
    if isinstance(event, list):
        # 如果是事件列表，转换为张量
        x_coords = torch.tensor([e['x'] for e in event], dtype=torch.long, device=device)
        y_coords = torch.tensor([e['y'] for e in event], dtype=torch.long, device=device)
        p_vals = torch.tensor([e['p'] for e in event], dtype=torch.float32, device=device)
        if num_frame > 1:
            t_vals = torch.tensor([e['t'] for e in event], dtype=torch.float32, device=device)
    else:
        # 如果已经是张量格式
        x_coords = torch.from_numpy(event['x'].astype(np.int32)).to(device)
        y_coords = torch.from_numpy(event['y'].astype(np.int64)).to(device)
        p_vals = torch.from_numpy(event['p'].astype(np.float32)).to(device)
        if num_frame > 1:
            t_vals = torch.from_numpy(event['t'].astype(np.float32)).to(device)
    
    if num_frame == 1:
        # 单帧处理
        # 初始化图像为128
        image = torch.full((h, w), 128, dtype=torch.float32, device=device)
        
        # 创建正事件和负事件的掩码
        pos_mask = (p_vals == 1)
        neg_mask = (p_vals <1) # ==-1 有可能不是-1而是0！！！
        
        # 使用bincount进行累加
        pos_counts = torch.bincount(y_coords[pos_mask] * w + x_coords[pos_mask], minlength=h*w).reshape(h, w)
        neg_counts = torch.bincount(y_coords[neg_mask] * w + x_coords[neg_mask], minlength=h*w).reshape(h, w)
        
        # 计算最终图像
        image = image + pos_counts - neg_counts
        
        # 限制在[0, 255]范围内
        # image = torch.clamp(image, 0, 255)


        image = torch.where(image >= 128 + threshold, torch.tensor(200.0).to(device), image)
        image = torch.where(image <= 128 - threshold, torch.tensor(100.0).to(device), image)
        image = torch.where((image !=200) & (image!=100), torch.tensor(0.0).to(device), image)  

        # HW ==》 HW3
        image = image.unsqueeze(2).repeat(1, 1, channel)

        
        # 转换为uint8并返回
        return image.byte().cpu().numpy()
    
    else:
        # 多帧处理
        # 获取时间范围
        t_min, t_max = t_vals.min(), t_vals.max()
        
        # 创建时间分箱
        t_bins = torch.linspace(t_min, t_max, num_frame + 1, device=device)
        
        # 为每个帧分配事件
        frame_indices = torch.bucketize(t_vals, t_bins, right=True) - 1
        frame_indices = torch.clamp(frame_indices, 0, num_frame - 1)


        if needed == -1: # just need last one 

            # 初始化图像为128
            image = torch.full((h, w), 128, dtype=torch.float32, device=device)
            
            # 获取当前帧的事件
            frame_mask = (frame_indices == num_frame-1)

            frame_x = x_coords[frame_mask]
            frame_y = y_coords[frame_mask]
            frame_p = p_vals[frame_mask]

            # 正负事件掩码
            pos_mask = (frame_p == 1)
            neg_mask = (frame_p <  1)

            # 使用bincount进行累加

            if len(frame_x) > 0:
                pos_counts = torch.bincount(frame_y[pos_mask] * w + frame_x[pos_mask], minlength=h*w).reshape(h, w)
                neg_counts = torch.bincount(frame_y[neg_mask] * w + frame_x[neg_mask], minlength=h*w).reshape(h, w)
                
                image = image + pos_counts - neg_counts

            image = torch.where(image >= 128 + threshold, torch.tensor(200.0).to(device), image)
            image = torch.where(image <= 128 - threshold, torch.tensor(100.0).to(device), image)
            image = torch.where((image !=200) & (image!=100), torch.tensor(0.0).to(device), image)

            image = image.unsqueeze(2).repeat(1, 1, channel)  
        
            # 转换为uint8并返回
            return image.byte().cpu().numpy()

        else:

            images = []

            for i in range(num_frame):
                # 初始化图像为128
                image = torch.full((h, w), 128, dtype=torch.float32, device=device)
                
                # 获取当前帧的事件
                frame_mask = (frame_indices == i)
                frame_x = x_coords[frame_mask]
                frame_y = y_coords[frame_mask]
                frame_p = p_vals[frame_mask]
                
                # 正负事件掩码
                pos_mask = (frame_p == 1)
                neg_mask = (frame_p <  1)
                
                # 使用bincount进行累加
                if len(frame_x) > 0:
                    pos_counts = torch.bincount(frame_y[pos_mask] * w + frame_x[pos_mask], minlength=h*w).reshape(h, w)
                    neg_counts = torch.bincount(frame_y[neg_mask] * w + frame_x[neg_mask], minlength=h*w).reshape(h, w)
                    
                    image = image + pos_counts - neg_counts

                image = torch.where(image >= 128 + threshold, torch.tensor(200.0).to(device), image)
                image = torch.where(image <= 128 - threshold, torch.tensor(100.0).to(device), image)
                image = torch.where((image !=200) & (image!=100), torch.tensor(0.0).to(device), image)  
                
                image = image.unsqueeze(2).repeat(1, 1, channel)

                images.append(image.byte().cpu().numpy())

            return images


if __name__ == '__main__':
    one = torch.ones(1,10)
    zero = torch.zeros(1,10)

    all = torch.cat([one,zero],dim=0)
    all_2 = torch.stack([one,zero], dim=0)

    print (all.shape) #torch.Size([2, 10])

    print (all_2.shape)# torch.Size([2, 1, 10])